import cv2
import numpy as np
import random
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont

from .classes import CLASSES


class Drawer:

    @staticmethod
    def init(config):
        config.box_colors = tuple(tuple(random.randint(0, 255) for _ in range(3)) for _ in range(len(CLASSES)))
    
    def __init__(self, config):
        self._font = ImageFont.truetype(str(Path(__file__).parent / 'simhei.ttf'), int(40 * config.stream_w / 1920))
        self._colors = config.box_colors
        self._draw_results = config.enable_draw_video_box
        
        # DBG
        # self._highlight_pts = []
        # for poly in config.ai_region:
        #     highlight_pt = np.array(poly['region']).reshape(-1, 2)
        #     highlight_pt[:, 0] *= config.stream_w
        #     highlight_pt[:, 1] *= config.stream_h
        #     highlight_pt = highlight_pt.astype(np.int32)
        #     self._highlight_pts.append(highlight_pt)

        # self._highlight_mask = cv2.fillPoly(
        #     np.zeros((config.stream_h, config.stream_w, 3), dtype=np.uint8), 
        #     self._highlight_pts, 
        #     (0, 165, 255)
        # )

    # 用Pillow画图，但是耗时比cv搞一个量级，不然没法显示中文
    # 返回一个画好的新图
    def draw(self, image, results):
        # image = image.copy()
        # cv2.addWeighted(image, 1, self._highlight_mask, 0.2, 0, image) # DBG
        # cv2.polylines(image, self._highlight_pts, True, (255, 255, 0), 1) # DBG
        res_image = image
        if results is not None and self._draw_results:
            image = Image.fromarray(image)
            draw = ImageDraw.Draw(image)
            for box, cl, score in results:
                text = CLASSES[cl]
                draw.rectangle(box, outline=self._colors[cl], width=2)
                label_box = draw.textbbox((0,0), text, self._font)
                y = box[1] - label_box[3]
                if y < 0:
                    y = box[3] + 1
                draw.text((box[0], y), text, self._colors[cl], self._font)
            res_image = np.asarray(image)
        return res_image

    # 直接在传入图像上画
    # def draw_results(image, results):
    #     for box, cl, score in results:
    #         top, left, right, bottom = box

    #         cv2.rectangle(image, (top, left), (right, bottom), colors[cl], 2)
    #         cv2.putText(image, 
    #                     CLASSES[cl],
    #                     (top, left - 6),
    #                     cv2.FONT_HERSHEY_COMPLEX,
    #                     0.6, colors[cl], 2)
